#Todo: 完善单批次训练的函数
# Input: 每个批次的data数据
# output: 预测准确度


def train_one_epoch(model, data, optimizer, criterion, task_type, device, train=True):
    data = data.to(device)
    out = model(data.x, data.edge_index)
    loss = criterion(out, data.y)

    if train:
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    metrics = {'loss': loss.item()}

    if task_type == 'classification':
        pred = out.argmax(dim=1)
        acc = (pred == data.y).sum().item() / data.y.size(0)
        metrics['acc'] = acc

    return metrics


